//=============================================================================
//
//   Copyright (C) 2020 All rights reserved.
//
//   Filename:		tensor.rs
//
//   Author: Wang Zhecheng - wangzhecheng@yeah.net
//
//   Date: 2020-10-03 02:15
//
//   Last Editors: Wang Zhecheng - wangzhecheng@yeah.net
//
//   Last modified:	2020-11-06 18:27
//
//   Description:
//
//=============================================================================
use crate::consts::{FRAC_1_SQRT_3, PI, ROOT2, ROOT3};
use crate::dimension::{DimType, PlaneStrain, PlaneStress, ThreeDimension};
use crate::Real;
use nalgebra::allocator::Allocator;
use nalgebra::dimension::{Dim, U1, U3};
use nalgebra::{
    DefaultAllocator, Matrix3, MatrixN, MatrixSliceMutN, MatrixSliceN, Vector6, VectorN,
};
use std::fmt;
use std::marker::PhantomData;
use std::ops::{Index, IndexMut};

pub trait TensorType: Default + fmt::Display + Clone + Copy {}
#[derive(Default, Debug, Copy, Clone, PartialEq)]
pub struct StressTag {}
impl TensorType for StressTag {}
impl fmt::Display for StressTag {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "stress")
    }
}
#[derive(Default, Debug, Copy, Clone, PartialEq)]
pub struct StrainTag {}
impl TensorType for StrainTag {}
impl fmt::Display for StrainTag {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "strain")
    }
}

#[derive(Debug, Copy, Clone, PartialEq)]
pub struct Tensor<D, T> {
    value: Matrix3<Real>,
    _marker1: PhantomData<D>,
    _marker2: PhantomData<T>,
}

pub type Stress<D> = Tensor<D, StressTag>;
pub type Strain<D> = Tensor<D, StrainTag>;

impl<D: DimType, T: TensorType> Tensor<D, T> {
    #[inline]
    pub fn value(&self) -> &Matrix3<Real> {
        &self.value
    }
    #[inline]
    pub fn value_mut(&mut self) -> &mut Matrix3<Real> {
        &mut self.value
    }
    #[inline]
    pub fn tensor(&self) -> MatrixSliceN<Real, D::Dim, U1, U3> {
        self.value().fixed_slice::<D::Dim, D::Dim>(0, 0)
    }
    #[inline]
    pub fn tensor_mut(&mut self) -> MatrixSliceMutN<Real, D::Dim, U1, U3> {
        self.value_mut().fixed_slice_mut::<D::Dim, D::Dim>(0, 0)
    }
    #[inline]
    pub fn set_value(&mut self, rhs: &MatrixN<Real, D::Dim>) -> Self
    where
        DefaultAllocator: Allocator<Real, D::Dim, D::Dim>,
    {
        self.value
            .fixed_slice_mut::<D::Dim, D::Dim>(0, 0)
            .copy_from(rhs);
        *self
    }
    #[inline]
    pub fn reset(&mut self) -> Self {
        self.value_mut().fill(0 as Real);
        *self
    }
    #[inline]
    pub fn from_voigt(rhs: VectorN<Real, D::DimVoigt>) -> Self
    where
        DefaultAllocator: Allocator<Real, D::Dim, D::Dim> + Allocator<Real, D::DimVoigt>,
    {
        Self {
            value: D::from_voigt(&rhs).fixed_resize::<U3, U3>(0 as Real),
            _marker1: PhantomData::<D> {},
            _marker2: PhantomData::<T> {},
        }
    }
    #[inline]
    pub fn to_voigt(&self) -> VectorN<Real, D::DimVoigt>
    where
        DefaultAllocator: Allocator<Real, D::Dim, D::Dim> + Allocator<Real, D::DimVoigt>,
    {
        D::to_voigt(&self.tensor().clone_owned())
    }
    #[inline]
    pub fn dm(&self) -> Real {
        self.i1() / 3 as Real
    }
    #[inline]
    pub fn s11(&self) -> Real {
        self[(0, 0)]
    }
    #[inline]
    pub fn s22(&self) -> Real {
        self[(1, 1)]
    }
    #[inline]
    pub fn s33(&self) -> Real {
        self[(2, 2)]
    }
    /// # fn i1(&self) -> Real
    /// `i1` is a function  that calculate tensor's invariant 1
    /// $$ I2 = d11 + d22 + d33 $$
    #[inline]
    pub fn i1(&self) -> Real {
        self.value().trace()
    }
    /// # fn i2(&self) -> Real
    /// `i2` is a function  that calculate tensor's invariant 2
    /// $$ I2=d11*d22+d22*d33+d11*d33-d12^2-d13^2-d23^2 $$
    #[inline]
    pub fn i2(&self) -> Real {
        self[(0, 0)] * self[(1, 1)] + self[(1, 1)] * self[(2, 2)] + self[(0, 0)] * self[(2, 2)]
            - self[(0, 1)] * self[(0, 1)]
            - self[(1, 2)] * self[(1, 2)]
            - self[(0, 2)] * self[(0, 2)]
    }
    /// # fn i3(&self) -> Real
    /// `i3` is a function  that calculate tensor's invariant 3
    /// $ I3= d11*d22*d33 + 2d12*d13*d23 - d11*d23*d23 - d22*d13*d13 - d33*d12*d12 $
    #[inline]
    pub fn i3(&self) -> Real {
        self.value().determinant()
    }
    /// get the derivatives of I1 with respect to the component
    #[inline]
    pub fn di1ds(&self) -> Vector6<Real> {
        Vector6::<Real>::new(1., 1., 1., 0., 0., 0.)
    }
    /// get the derivatives of I2 with respect to the component
    #[inline]
    pub fn di2ds(&self) -> Vector6<Real> {
        Vector6::<Real>::new(
            self[(1, 1)] + self[(2, 2)],
            self[(0, 0)] + self[(2, 2)],
            self[(0, 0)] + self[(1, 1)],
            -2. * self[(0, 1)],
            -2. * self[(0, 2)],
            -2. * self[(1, 2)],
        )
    }
    /// get the derivatives of I3 with respect to the component
    #[inline]
    pub fn di3ds(&self) -> Vector6<Real> {
        Vector6::<Real>::new(
            self[(1, 1)] * self[(2, 2)] - self[(1, 2)] * self[(1, 2)],
            self[(0, 0)] * self[(2, 2)] - self[(0, 2)] * self[(0, 2)],
            self[(0, 0)] * self[(1, 1)] - self[(0, 1)] * self[(0, 1)],
            2. * self[(0, 2)] * self[(1, 2)] - 2. * self[(0, 1)] * self[(2, 2)],
            2. * self[(0, 1)] * self[(1, 2)] - 2. * self[(0, 2)] * self[(1, 1)],
            2. * self[(0, 1)] * self[(0, 2)] - 2. * self[(1, 2)] * self[(0, 0)],
        )
    }
    /// get deviatoric tensor packetp
    #[inline]
    pub fn deviatoric(&self) -> MatrixN<Real, D::Dim>
    where
        DefaultAllocator: Allocator<Real, D::Dim, D::Dim>,
    {
        self.tensor() - MatrixN::<Real, D::Dim>::identity() * self.dm() / 3.
    }
    /// get deviatoric tensor invariant 1
    #[inline]
    pub fn j1(&self) -> Real {
        0 as Real
    }
    /// get deviatoric tensor invariant 2
    /// also write as 1/6*[(d11-d22)^2+(d22-d33)^2+(d11-d33)^2]+d12^2+d13^2+d23^2
    #[inline]
    pub fn j2(&self) -> Real {
        let i1 = self.i1();
        let i2 = self.i2();
        i1 * i1 / 3. - i2
    }
    /// get deviatoric tensor invariant 3
    #[inline]
    pub fn j3(&self) -> Real {
        let i1 = self.i1();
        let i2 = self.i2();
        let i3 = self.i3();
        2. * i1 * i1 * i1 / 27. - i1 * i2 / 3. + i3
    }
    /// get the derivatives of T with respect to the component
    #[inline]
    pub fn dpds(&self) -> Vector6<Real> {
        self.di1ds() / 3.
    }
    /// get the derivatives of J2 with respect to the component
    #[inline]
    pub fn dj2ds(&self) -> Vector6<Real> {
        Vector6::<Real>::new(
            self.s11(),
            self.s22(),
            self.s33(),
            2. * self[(0, 1)],
            2. * self[(0, 2)],
            2. * self[(1, 2)],
        )
    }
    /// get the derivatives of J3 with respect to the component
    #[inline]
    pub fn dj3ds(&self) -> Vector6<Real>
    where
        DefaultAllocator: Allocator<Real, D::Dim, D::Dim>,
    {
        let s = self.deviatoric();
        let dj3ds = &s * &s - MatrixN::<Real, D::Dim>::identity() * self.j2() * 2. / 3.;
        Vector6::<Real>::new(
            dj3ds[(0, 0)],
            dj3ds[(1, 1)],
            dj3ds[(2, 2)],
            dj3ds[(0, 1)] + dj3ds[(1, 0)],
            dj3ds[(0, 2)] + dj3ds[(2, 0)],
            dj3ds[(1, 2)] + dj3ds[(2, 1)],
        )
    }
    /// get /sigma_{Octahedral}
    #[inline]
    pub fn soct(&self) -> Real {
        (2. as Real / 3. as Real * self.j2()).sqrt()
    }
    /// get mean stress
    #[inline]
    pub fn p(&self) -> Real {
        self.dm() / 3.
    }
    /// get deviatoric stress
    #[inline]
    pub fn q(&self) -> Real {
        (3. as Real * self.j2()).sqrt()
    }
    /// get Von Mises stress
    #[inline]
    pub fn mises(&self) -> Real {
        (3. as Real * self.j2()).sqrt()
    }
    /// effective stress
    #[inline]
    pub fn sbar(&self) -> Real {
        (3. as Real * self.j2()).sqrt()
    }
    /// get the derivatives of q with respect to the component
    #[inline]
    pub fn dqds(&self) -> Vector6<Real> {
        if self.j2() == 0. {
            Vector6::<Real>::zeros()
        } else {
            self.dj2ds() * (3. as Real / self.j2()).sqrt() / 2. as Real
        }
    }
    /// get the derivatives of Mises with respect to the component
    #[inline]
    pub fn dmisesds(&self) -> Vector6<Real> {
        if self.j2() == 0. {
            Vector6::<Real>::zeros()
        } else {
            (3. as Real / self.j2()).sqrt() * self.dj2ds() / 2.
        }
    }
    /// get the derivatives of rho with respect to the component
    #[inline]
    pub fn drhods(&self) -> Vector6<Real> {
        if self.j2() == 0. {
            Vector6::<Real>::zeros()
        } else {
            self.dj2ds() * (2. as Real / self.j2()).sqrt() / 2.
        }
    }
    /// get lode theta. anti-clockwise. start from x-axis direction at PI plane domain: (-PI/6,PI/6)
    #[inline]
    pub fn lode2(&self) -> Real {
        if self.j2() == 0. {
            0.
        } else {
            1. as Real / 3. * (-(3. as Real / self.j2()).powf(1.5) * self.j3() * 0.5).asin()
        }
    }
    /// get  Haigh-Westergaard coordinates Xi
    #[inline]
    pub fn xi(&self) -> Real {
        self.i1() * FRAC_1_SQRT_3
    }
    /// get  Haigh-Westergaard coordinates rho
    #[inline]
    pub fn rho(&self) -> Real {
        (2. as Real * self.j2()).sqrt()
    }
    ///  get Haigh-Westergaard coordinates lode angle. clockwise. start from 150 degree direction at PI plane domain: (-PI/6,PI/6)
    #[inline]
    pub fn lode(&self) -> Real {
        if self.j2() == 0. {
            0.
        } else {
            1. as Real / 3. * ((3. as Real / self.j2()).powf(1.5) * self.j3() / 2.).acos()
        }
    }
    /// get the derivatives of lode lode with respect to sigma
    #[inline]
    pub fn dlodeds(&self) -> Vector6<Real>
    where
        DefaultAllocator: Allocator<Real, D::Dim, D::Dim>,
    {
        let lode = self.lode();
        let j2 = self.j2();
        let j3 = self.j3();

        ROOT3 / 2. / (3. * lode).sin() / j2.powf(1.5)
            * (1.5 * j3 / j2 * self.dj2ds() - self.dj3ds())
    }

    /// get maximum principal stress
    #[inline]
    pub fn pmax(&self) -> Real {
        self.xi() * FRAC_1_SQRT_3 + ROOT2 * FRAC_1_SQRT_3 * self.rho() * self.lode().cos()
    }
    /// get mid principal stress
    #[inline]
    pub fn pmid(&self) -> Real {
        self.xi() * FRAC_1_SQRT_3
            + ROOT2 * FRAC_1_SQRT_3 * self.rho() * (self.lode() - 2. / 3. * PI).cos()
    }
    /// get minimum principal stress
    #[inline]
    pub fn pmin(&self) -> Real {
        self.xi() * FRAC_1_SQRT_3
            + ROOT2 * FRAC_1_SQRT_3 * self.rho() * (self.lode() + 2. / 3. * PI).cos()
    }
}

///
/// Index
///
impl<D: DimType, T: TensorType> Index<(usize, usize)> for Tensor<D, T> {
    type Output = Real;
    #[inline]
    fn index(&self, index: (usize, usize)) -> &Self::Output {
        let dim = match D::Dim::try_to_usize() {
            None => panic!("D::Dim is not known"),
            Some(i) => i,
        };
        match index {
            (x, y) if x < dim && y < dim => &self.value()[(x, y)],
            _ => &(0 as Real),
        }
    }
}
impl<D: DimType, T: TensorType> IndexMut<(usize, usize)> for Tensor<D, T> {
    #[inline]
    fn index_mut(&mut self, index: (usize, usize)) -> &mut Self::Output {
        unsafe { self.value_mut().get_unchecked_mut(index) }
    }
}

///
/// Display
///
impl<D: DimType, T: TensorType> fmt::Display for Tensor<D, T> {
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}:{}", T::default(), self.value())
    }
}

///
/// Default
///
impl<D: DimType, T: TensorType> Default for Tensor<D, T> {
    #[inline]
    fn default() -> Self {
        Self {
            value: Matrix3::<Real>::zeros(),
            _marker1: PhantomData::<D> {},
            _marker2: PhantomData::<T> {},
        }
    }
}

///
/// constructor new
///
macro_rules! constructors_impl(
    ($($DimType: ty, $($args: ident:($irow: expr,$icol: expr)),*);* $(;)*) => {$(
        impl Tensor<$DimType,StrainTag>
        {
            /// Initializes this matrix from its components.
            #[inline]
            pub fn new($($args: Real),*) -> Self {
                unsafe {
                    let mut res = Self::default();
                    $( *res.value.get_unchecked_mut(($irow, $icol)) = $args; )*
                    res
                }
            }
        }
        impl Tensor<$DimType,StressTag>
        {
            /// Initializes this matrix from its components.
            #[inline]
            pub fn new($($args: Real),*) -> Self {
                unsafe {
                    let mut res = Self::default();
                    $( *res.value.get_unchecked_mut(($irow, $icol)) = $args; )*
                    res
                }
            }
        }
    )*}
);

constructors_impl!(

PlaneStrain,
        m11:(0,0), m12:(0,1),
        m21:(1,0), m22:(1,1);
PlaneStress,
        m11:(0,0), m12:(0,1),
        m21:(1,0), m22:(1,1);
ThreeDimension,
        m11:(0,0), m12:(0,1), m13:(0,2),
        m21:(1,0), m22:(1,1), m23:(1,2),
        m31:(2,0), m32:(2,1), m33:(2,2);
        );

#[test]
fn constructor() {
    use crate::dimension::{PlaneStress, ThreeDimension};
    use nalgebra::Matrix2;
    let mut t1 = Stress::<PlaneStress>::new(1f32, 3f32, 3f32, 2f32);
    assert_eq!(t1[(0, 0)], 1f32);
    assert_eq!(t1[(0, 1)], 3f32);
    assert_eq!(t1[(1, 0)], 3f32);
    assert_eq!(t1[(1, 1)], 2f32);
    let v = Matrix2::<Real>::new(3f32, 8f32, 0f32, 6f32);
    t1.set_value(&v);
    assert_eq!(t1[(0, 0)], 3f32);
    assert_eq!(t1[(0, 1)], 8f32);
    assert_eq!(t1[(1, 0)], 0f32);
    assert_eq!(t1[(1, 1)], 6f32);
    let t2 = Stress::<ThreeDimension>::new(
        1f32, 4f32, 5f32, //
        4f32, 2f32, 6f32, //
        5f32, 6f32, 3f32,
    );
    assert_eq!(t2[(0, 0)], 1f32);
    assert_eq!(t2[(0, 1)], 4f32);
    assert_eq!(t2[(0, 2)], 5f32);
    assert_eq!(t2[(1, 0)], 4f32);
    assert_eq!(t2[(1, 1)], 2f32);
    assert_eq!(t2[(1, 2)], 6f32);
    assert_eq!(t2[(2, 0)], 5f32);
    assert_eq!(t2[(2, 1)], 6f32);
    assert_eq!(t2[(2, 2)], 3f32);
}
#[test]
fn voigt() {
    use crate::dimension::{PlaneStress, ThreeDimension};
    use nalgebra::Vector3;
    let t1 = Stress::<PlaneStress>::from_voigt(Vector3::<Real>::new(1f32, 2f32, 3f32));
    assert_eq!(t1[(0, 0)], 1f32);
    assert_eq!(t1[(0, 1)], 3f32);
    assert_eq!(t1[(1, 0)], 3f32);
    assert_eq!(t1[(1, 1)], 2f32);
    let v1 = t1.to_voigt();
    assert_eq!(v1[0], 1f32);
    assert_eq!(v1[1], 2f32);
    assert_eq!(v1[2], 3f32);
    let t2 = Stress::<ThreeDimension>::from_voigt(Vector6::<Real>::new(
        1f32, 2f32, 3f32, 4f32, 5f32, 6f32,
    ));
    assert_eq!(t2[(0, 0)], 1f32);
    assert_eq!(t2[(0, 1)], 4f32);
    assert_eq!(t2[(0, 2)], 5f32);
    assert_eq!(t2[(1, 0)], 4f32);
    assert_eq!(t2[(1, 1)], 2f32);
    assert_eq!(t2[(1, 2)], 6f32);
    assert_eq!(t2[(2, 0)], 5f32);
    assert_eq!(t2[(2, 1)], 6f32);
    assert_eq!(t2[(2, 2)], 3f32);
    let v2 = t2.to_voigt();
    assert_eq!(v2[0], 1f32);
    assert_eq!(v2[1], 2f32);
    assert_eq!(v2[2], 3f32);
    assert_eq!(v2[3], 4f32);
    assert_eq!(v2[4], 5f32);
    assert_eq!(v2[5], 6f32);
}
#[test]
fn tensor2d() {
    use approx::assert_relative_eq;
    let t1_2d = Stress::<PlaneStress>::new(
        80., -60., //
        -60., -40.,
    );
    assert_relative_eq!(t1_2d.pmax(), 104.8, max_relative = 1e-3);
    assert_relative_eq!(t1_2d.pmin(), -64.9, max_relative = 1e-3);
    let t2_2d = Stress::<PlaneStress>::new(
        0., 20., //
        20., -80.,
    );
    assert_relative_eq!(t2_2d.pmax(), 4.721, max_relative = 1e-3);
    assert_relative_eq!(t2_2d.pmin(), -84.721, max_relative = 1e-3);
}
#[test]
fn tensor3d() {
    use approx::assert_relative_eq;
    let t1_3d = Stress::<ThreeDimension>::new(
        5., 0., -5., //
        0., -5., 0., //
        -5., 0., 5.,
    );
    assert_relative_eq!(t1_3d.i1(), 5., max_relative = 1e-3);
    assert_relative_eq!(t1_3d.i2(), -50., max_relative = 1e-3);
    assert_relative_eq!(t1_3d.i3(), 0., epsilon = 1e-3, max_relative = 1e-3);
    assert_relative_eq!(t1_3d.pmax(), 10., max_relative = 1e-3);
    assert_relative_eq!(t1_3d.pmid(), 0., epsilon = 1e-3, max_relative = 1e-3);
    assert_relative_eq!(t1_3d.pmin(), -5., max_relative = 1e-3);
    assert_relative_eq!(t1_3d.i1() / 3., 5. / 3., max_relative = 1e-3);
    assert_relative_eq!(
        (t1_3d.j2() * 2. / 3.).sqrt(),
        5. / 3. * 14f32.sqrt(),
        max_relative = 1e-3
    );
    let t2_3d = Stress::<ThreeDimension>::new(
        0., -1., 1., //
        -1., 0., -1., //
        1., -1., 0.,
    );
    assert_relative_eq!(t2_3d.pmax(), 2., max_relative = 1e-3);
    assert_relative_eq!(t2_3d.pmid(), -1., max_relative = 1e-3);
    assert_relative_eq!(t2_3d.pmin(), -1., max_relative = 1e-3);
    assert_relative_eq!(t2_3d.j2(), 3., max_relative = 1e-3);
    assert_relative_eq!(t2_3d.lode2(), -PI / 6., max_relative = 1e-3);
    let t3_3d = Stress::<ThreeDimension>::new(
        7., 3., 0., //
        3., 7., 4., //
        0., 4., 7.,
    );
    assert_relative_eq!(t3_3d.pmax(), 12., max_relative = 1e-3);
    assert_relative_eq!(t3_3d.pmid(), 7., max_relative = 1e-3);
    assert_relative_eq!(t3_3d.pmin(), 2., max_relative = 1e-3);
    let t4_3d = Stress::<ThreeDimension>::new(
        -13., -1., -4., //
        -1., -13., -4., //
        -4., -4., -10.,
    );
    assert_relative_eq!(t4_3d.pmax(), -6., max_relative = 1e-3);
    assert_relative_eq!(t4_3d.pmid(), -12., max_relative = 1e-3);
    assert_relative_eq!(t4_3d.pmin(), -18., max_relative = 1e-3);
    assert_relative_eq!(t4_3d.lode2(), 0., epsilon = 1e-3, max_relative = 1e-3);
    assert_relative_eq!(t4_3d.j2(), 36., max_relative = 1e-3);
}
#[test]
fn display() {
    let t1_3d = Stress::<ThreeDimension>::new(
        5., 0., -5., //
        0., -5., 0., //
        -5., 0., 5.,
    );
    println!("{}", t1_3d);
}
